import time

import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS
from langchain.document_loaders import TextLoader
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from langchain.prompts import PromptTemplate
from request import generate_text
from es_search import search_and_save_results
from concurrent.futures import ThreadPoolExecutor
import os
from datetime import datetime

# 文件路径和模型路径
embedding_path = '/data/model/zpoint_large_embedding_zh'

# 创建线程池执行器
executor = ThreadPoolExecutor(max_workers=10)  # 根据系统调整线程数量

def task(name):
    print(f"Task {name} started")
    time.sleep(5)
    return f"ok"


class MyClassModel(BaseModel):
    question: str


def create_docsearch_index(file_path, embedding_path, query):
    try:
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件未找到: {file_path}")

        loader = TextLoader(file_path)
        documents = loader.load()

        text_splitter = RecursiveCharacterTextSplitter(chunk_size=0, chunk_overlap=0, separators=["-----"])
        texts = text_splitter.split_documents(documents)

        embeddings = HuggingFaceEmbeddings(model_name=embedding_path)
        docsearch = FAISS.from_documents(texts, embeddings)
        results = docsearch.similarity_search(query, k=10)
        return results[0].page_content
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

def prom(content, question, save_path):
    context = content
    prompt_template = (
        "{question}\n"
        "背景知识：\n{context}")
    prompt_text = prompt_template.format(context=context, question=question)
    PROMPT = PromptTemplate(template=prompt_text, input_variables=["context", "question"])
    os.remove(save_path)
    return PROMPT.template

def result(question, startYear, endYear):
    try:
        print(question)
        t1 = datetime.now()
        save_path = search_and_save_results(question, startYear, endYear)
        t2 = datetime.now()
        second = (t1 - t2).total_seconds()
        print(second)
        background = create_docsearch_index(save_path, embedding_path, question)
        t3 = datetime.now()
        second = (t3 - t2).total_seconds()
        print(second)

        prompt = prom(background, question, save_path)

        # re = generate_text(prompt)
        re = "ok"
        return re
    except Exception as e:
        return {"error": str(e)}

# 定义FastAPI应用
app = FastAPI()

# 定义请求数据模型
class Query(BaseModel):
    question: str
    startYear: Optional[int] = None
    endYear: Optional[int] = None

@app.post("/query")
async def query_endpoint(query: Query):
    start_time = time.time()
    try:
        response_future = asyncio.get_event_loop().run_in_executor(executor, result, query.question, query.startYear, query.endYear)
        response = await response_future  # 等待结果
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    end_time = time.time()
    elapsed_time = end_time - start_time
    print(f"处理时间: {elapsed_time} 秒")
    return response


@app.post("/test")
def create_myclass(myclass: MyClassModel):
    t1 = datetime.now()

    question = myclass.question
    print(question)

    future = executor.submit(task, question)
    # 获取任务执行结果
    print(future.result())

    result = future.result()

    t2 = datetime.now()
    second = (t1 - t2).total_seconds()

    print(f"耗时{second}")
    return {"answer": result}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=19302)
